/*
 * UserNet  --  TCP/IP implementation in userspace
 * Copyright (C) 2010  Robert Millan
 * Copyright (C) 2010  Felip Moll
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <assert.h>
#include <pthread.h>
#include <stdio.h>
#include <stdint.h>
#include <errno.h>
#include <fcntl.h>		/* open */
#include <unistd.h>
#include <sys/socket.h>

#include "unet/unet.h"
#include "unet/kernel.h"
#include "unet/misc.h"
#include "unet/ip.h"

enum sock_types
  {
    SOCK_TYPE_UNDEFINED,
    SOCK_TYPE_LISTEN,
    SOCK_TYPE_CONNECT,
  };

struct unet_socket
{
  int fd;	/* Dummy file descriptor (see below).  */
  enum sock_types sock_type;
  enum ip_protocols protocol;
  uint16_t port;
  struct in_addr addr;
  void *state;
  struct unet_socket *next;
};

static struct unet_socket *socket_list[16];
static int socket_count = 0;
static pthread_mutex_t socket_list_mutex = PTHREAD_MUTEX_INITIALIZER;

static struct unet_socket *
create_socket (enum ip_protocols protocol)
{
  struct unet_socket *s;

  s = unet_zalloc (sizeof (*s));
  if (! s)
    goto fail0;

  /* This is a bit tricky... this file descriptor isn't actually
     going to be used (since we intercept syscalls), but we need
     it to identify the connection later.  We can't assign an
     arbitrary number because we need to garantee there won't be
     a collision with those assigned by the kernel.  */
  s->fd = open ("/dev/null", O_RDONLY);
  if (s->fd == -1)
    goto fail1;

  s->protocol = protocol;
  s->sock_type = SOCK_TYPE_UNDEFINED;

  pthread_mutex_lock (& socket_list_mutex);
  s->next = HASH_HEAD (socket_list, s->fd);
  HASH_HEAD (socket_list, s->fd) = s;
  pthread_mutex_unlock (& socket_list_mutex);

  assert (socket_count >= 0);

  if (socket_count == 0)
    {
      if (unet_init () == -1)
	goto fail2;
    }

  socket_count++;

  return s;

 fail2:
  real_close (s->fd);
 fail1:
  free (s);
 fail0:
  return NULL;
}

static struct unet_socket *
find_socket (int fd)
{
  struct unet_socket *s;

  for (s = HASH_HEAD (socket_list, fd); s; s = s->next)
    if (s->fd == fd)
      return s;

  return NULL;
}

int
socket (int namespace, int style, int protocol)
{
  struct unet_socket *s;

  if (namespace != AF_INET)
    goto not_for_us;

  if (style == SOCK_STREAM && (protocol == IPPROTO_TCP || protocol == 0))
    s = create_socket (UNET_TCP);
  else if (style == SOCK_DGRAM && protocol == IPPROTO_UDP)
    s = create_socket (UNET_UDP);
  else
    goto not_for_us;

  if (! s)
    return -1;

  return s->fd;

 not_for_us:
  return real_socket (namespace, style, protocol);
}

int
setsockopt (int fd, int level, int optname, const void *optval, socklen_t optlen)
{
  struct unet_socket *s = find_socket (fd);
  if (! s)
    return real_setsockopt (fd, level, optname, optval, optlen);

  errno = 0;

  // FIXME: unimplemented
  errno = ENOSYS;

  return errno ? -1 : 0;
}

int
getsockname (int fd, struct sockaddr *addr, socklen_t *len)
{
  struct unet_socket *s = find_socket (fd);
  if (! s)
    return real_getsockname (fd, addr, len);

  errno = 0;

  *len = sizeof (s->addr);
  memcpy (addr, &s->addr, *len);

  return *len;
}

int
connect (int fd, const struct sockaddr *addr, socklen_t addr_len)
{
  const struct sockaddr_in *addr_in = (void *) addr;
  struct unet_socket *s = find_socket (fd);
  if (! s)
    return real_connect (fd, addr, addr_len);

  errno = 0;

  s->port = ntohs (addr_in->sin_port);
  s->addr = addr_in->sin_addr;

  s->state = unet_connect (s->protocol, s->addr, s->port);
  if (s->state == NULL)
    return -1;

  s->sock_type = SOCK_TYPE_CONNECT;

  return 0;
}

ssize_t
read (int fd, void *buf, size_t n)
{
  struct unet_socket *s = find_socket (fd);
  size_t read_sz = 0;

  if (! s)
    return real_read (fd, buf, n);

  errno = 0;

  read_sz = unet_recv (s->state, buf, n);

  return errno ? -1 : read_sz;
}

ssize_t
recv (int fd, void *buf, size_t n, int flags)
{
  struct unet_socket *s = find_socket (fd);
  if (! s)
    return real_recv (fd, buf, n, flags);

  return read (fd, buf, n);
}

ssize_t
recvfrom (int fd, void *buf, size_t n, int flags, struct sockaddr *addr, socklen_t *addr_len)
{
  struct unet_socket *s = find_socket (fd);
  if (! s)
    return real_recvfrom (fd, buf, n, flags, addr, addr_len);

  return read (fd, buf, n);
}

ssize_t
write (int fd, const void *buf, size_t n)
{
  struct unet_socket *s = find_socket (fd);

  if (! s)
    return real_write (fd, buf, n);

  errno = 0;

  unet_send (s->state, buf, n);

  return errno ? -1 : n;
}

int
bind (int fd, const struct sockaddr *addr, socklen_t addr_len)
{
  const struct sockaddr_in *addr_in = (void *) addr;
  struct unet_socket *s = find_socket (fd);
  if (! s)
    return real_bind (fd, addr, addr_len);

  errno = 0;

  s->port = ntohs (addr_in->sin_port);
  s->addr = addr_in->sin_addr;

  /* Listen with a zero-length queue.  */
  /*s->state = unet_listen (s->protocol, s->addr, s->port, 0);
  if (! s->state)
    return -1;
  */
  s->sock_type = SOCK_TYPE_LISTEN;

  return 0;
}

int
listen (int fd, int n)
{
  struct unet_socket *s = find_socket (fd);
  if (! s)
    return real_listen (fd, n);

  errno = 0;

  // unet_listen_close (s->state);
    s->state = unet_listen (s->protocol, s->addr, s->port, n);
  if (! s->state)
    return -1;

  return 0;
}

int
accept (int listenfd, struct sockaddr *addr, socklen_t *addr_len)
{
  struct unet_socket *s = find_socket (listenfd);
  struct unet_socket *accept_socket;
  void *state;
  if (! s)
    return real_accept (listenfd, addr, addr_len);

  errno = 0;

  state = unet_accept (s->state);
  if (! state)
    return -1;

  getsockname (listenfd, addr, addr_len);

  /* Create a new socket.  */
  accept_socket = create_socket (s->protocol);

  if (! accept_socket)
    return -1;

  accept_socket->port = s->port;
  accept_socket->addr = s->addr;

  accept_socket->state = state;

  accept_socket->sock_type = SOCK_TYPE_CONNECT;

  return accept_socket->fd;
}

int
close (int fd)
{
  struct unet_socket *s = find_socket (fd);
  if (! s)
    return real_close (fd);

  errno = 0;

  /* Ask lower layer to free resources.  */

  switch (s->sock_type)
    {
    case SOCK_TYPE_CONNECT:
      unet_conn_close (s->state);
      break;
    case SOCK_TYPE_LISTEN:
      unet_listen_close (s->state);
      break;
    case SOCK_TYPE_UNDEFINED:
      break;
    }

  /* Remove socket from list.  */
  pthread_mutex_lock (& socket_list_mutex);
  if (s == HASH_HEAD (socket_list, fd))
    {
      HASH_HEAD (socket_list, fd) = HASH_HEAD (socket_list, fd)->next;
    }
  else
    {
      struct unet_socket *i;
      for (i = HASH_HEAD (socket_list, fd); i; i = i->next)
	if (i->next == s)
	  {
	    i->next = i->next->next;
	    break;
	  }
    }
  pthread_mutex_unlock (& socket_list_mutex);
  free (s);

  /* Close the dummy socket (it was a /dev/null handle).  */
  real_close (fd);

  assert (socket_count >= 0);

  if (--socket_count == 0)
    unet_fini ();

  return 0;
}
